import os
import re
import json
import time
import logging
import asyncio
from typing import List
from collections import defaultdict
from fastapi import FastAPI

from openai import AsyncOpenAI
import uvicorn
from Agent.Agent_Memory import Agent_Memory
from Agent.Agent_Memory import agent_role_design
from Agent.Agent_Memory import remake_query_prompt_template
from  Log.LogManager import log_manager

# Configure logging
logger = log_manager.get_logger()

# ============================================================================
# Configuration Management
# ============================================================================
API_KEY = os.environ.get("ARK_API_KEY", "sk-4GBdGkeXISUn8sZl339Cy5RN2UjTb0Wqik6uyPWl6oOvM6QC")
BASE_URL = os.environ.get("ARK_API_BASE_URL", "http://ai.docker.tcl.com/huizhou/inference/usr-gqb23upd/inf-4fn1yn0y/v1")  # 例如 https://api.openai.com/v1
MODEL_NAME = os.environ.get("ARK_MODEL_NAME", "")

def validate_configuration():
    """Validate that required configuration values are set"""
    errors = []
    
    if not API_KEY or API_KEY == "sk-4GBdGkeXISUn8sZl339Cy5RN2UjTb0Wqik6uyPWl6oOvM6QC":
        errors.append("ARK_API_KEY is not properly configured")
    
    if not BASE_URL or BASE_URL == "http://ai.docker.tcl.com/huizhou/inference/usr-gqb23upd/inf-4fn1yn0y/v1":
        errors.append("ARK_API_BASE_URL is not properly configured")
    
    return errors

client = AsyncOpenAI(
    api_key=API_KEY,
    base_url=BASE_URL
)

# ============================================================================
# Agent Memory Management
# ============================================================================

# In-memory storage for user conversations
user_memories = defaultdict(Agent_Memory)
# Timestamps for cleanup
memory_timestamps = defaultdict(float)

def get_user_memory(user_identity: str) -> Agent_Memory:
    """Get or create memory for a user"""
    if user_identity not in user_memories:
        user_memories[user_identity] = Agent_Memory(user_identity)
        logger.info(f"Created new memory for user: {user_identity}")
    else:
        logger.debug(f"Retrieved existing memory for user: {user_identity}")
    
    memory_timestamps[user_identity] = time.time()
    return user_memories[user_identity]

def cleanup_old_memories(max_age_seconds: int = 3600):
    """Remove memories older than max_age_seconds"""
    current_time = time.time()
    expired_users = [
        user_id for user_id, timestamp in memory_timestamps.items()
        if current_time - timestamp > max_age_seconds
    ]
    
    if expired_users:
        logger.info(f"Cleaning up {len(expired_users)} expired user memories")
    
    for user_id in expired_users:
        if user_id in user_memories:
            del user_memories[user_id]
        if user_id in memory_timestamps:
            del memory_timestamps[user_id]
        
        logger.debug(f"Removed memory for user: {user_id}")


async def suggest_agent(memory: Agent_Memory):
    user_id = memory.user_identity
    history = memory.merge_memory()
    
    try:
        logger.info(f"Calling AI model for user: {user_id}")
        response = await client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": agent_role_design},
                {"role": "user", "content": remake_query_prompt_template.format(user_id, history)}
            ],
            extra_body={"enable_think": False},
            stream=False
        )

        content = response.choices[0].message.content
        

        # Pre-compile regex patterns for better performance
        think_tag_pattern = re.compile(r'\<think\>[\s\S]*?\<\/think\>')
        leading_json_pattern = re.compile(r'^(\n+)?```json\n')
        trailing_json_pattern = re.compile(r'\n```$')

        # 清理思维链
        content = think_tag_pattern.sub('', content).strip()
        content = leading_json_pattern.sub('', content, count=1)
        content = trailing_json_pattern.sub('', content, count=1)

        try:
            suggest_json = json.loads(content)
        except json.JSONDecodeError:
            logger.error(f"Failed to parse JSON response for user: {user_id}")
            suggest_json = [{"qid": "0", "content": "解析失败，原始输出：" + content}]
        return suggest_json
    except Exception as e:
        # Handle API call errors
        logger.error(f"API call failed for user {user_id}: {str(e)}")
        return [{"qid": "0", "content": f"API调用失败: {str(e)}"}]

# ============================================================================
# FastAPI Application
# ============================================================================
app = FastAPI()

from fastapi import HTTPException
from response.responseModel import ResponseModel
from response.HttpStatus import HttpStatus,Status
from entity.QueryRequest import QueryRequest

@app.post("/suggest", response_model=ResponseModel)
async def suggest_api(req: QueryRequest):
    start_time = time.time()  # 添加这行代码
    try:
        logger.info(f"Received request for user: {req.user_identity}")
        logger.debug(f"User question: {req.user_question}")
        
        # Get or create memory for the user
        memory = get_user_memory(req.user_identity)
        # Add the current query to memory
        memory.add_query(req.user_question)

        suggestions = await suggest_agent(memory)
        
        # Clean up old memories periodically
        # In a production system, this should be done in a background task
        cleanup_old_memories()
        
        logger.info(f"Returning {len(suggestions)} suggestions for user: {req.user_identity}")
        end_time = time.time()
        execution_time = end_time - start_time
        
        logger.info(f"suggest_agent function executed in {execution_time:.4f}  seconds \n User question: {req.user_question} \n Guess you Want to Ask : {suggestions}")
        #return {"recommendations": suggestions}
        return ResponseModel(msg=Status.SUCCESS,code=HttpStatus.OK,data=suggestions)
    except ValueError as e:
        logger.warning(f"Validation error for user {req.user_identity}: {str(e)}")
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Unexpected error for user {req.user_identity}: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal server error")

# ============================================================================
# Main Function (for testing)
# ============================================================================
async def test_suggest_agent():
    """Test the suggest_agent function with a sample memory"""
    # Create a test memory
    record = Agent_Memory()
    record.user_identity = "财务经理"

    user_question = "公司今年的现金流情况怎么样？"
    record.add_query(user_question)

    # Call the suggest_agent function
    suggestions = await suggest_agent(record)
    
    # Print results
    print("用户提问：", user_question)
    print("推荐的问题：", json.dumps(suggestions, ensure_ascii=False, indent=2))
    
    # Basic validation
    assert isinstance(suggestions, list), "Suggestions should be a list"
    if suggestions:
        first_suggestion = suggestions[0]
        assert "qid" in first_suggestion, "Each suggestion should have a 'qid' field"
        assert "content" in first_suggestion, "Each suggestion should have a 'content' field"
    
    print("Test passed!")
    return suggestions


if __name__ == '__main__':
    #test_suggest_agent()
    uvicorn.run(app, host="0.0.0.0", port=8000)